import paddle.fluid.layers as L
from paddle.fluid.dygraph import Layer, Sequential, Linear, load_dygraph

from helm.pretrained.layers import Conv2d, GlobalAvgPool, Pool2d, Identity, BN

__all__ = [
    "ResNeXt", "ResNeXt50_vd_64x4d", "ResNeXt101_vd_64x4d",
    "ResNeXt152_vd_64x4d", "ResNeXt50_vd_32x4d", "ResNeXt101_vd_32x4d",
    "ResNeXt152_vd_32x4d"
]

from torchvision.models import ResNet
class ResNeXt(Layer):

    def __init__(self, layers=50, groups=32, num_classes=1000):
        super().__init__()

        supported_layers = [50, 101, 152]
        assert layers in supported_layers, \
            "supported layers are {} but input layer is {}".format(supported_layers, layers)

        if layers == 50:
            depth = [3, 4, 6, 3]
        elif layers == 101:
            depth = [3, 4, 23, 3]
        elif layers == 152:
            depth = [3, 8, 36, 3]
        num_filters = [256, 512, 1024, 2048]
        cs = [64] + num_filters

        self.stem = Sequential(
            Conv2d(3, 32, 3, stride=2, bn=True, act='relu'),
            Conv2d(32, 32, 3, stride=1, bn=True, act='relu'),
            Conv2d(32, 64, 3, stride=1, bn=True, act='relu'),
        )

        self.pool = Pool2d(3, stride=2, type='max')

        self.layer1 = self._make_layer(cs[0], cs[1], depth[0], stride=1, groups=groups)
        self.layer2 = self._make_layer(cs[1], cs[2], depth[1], stride=2, groups=groups)
        self.layer3 = self._make_layer(cs[2], cs[3], depth[2], stride=2, groups=groups)
        self.layer4 = self._make_layer(cs[3], cs[4], depth[3], stride=2, groups=groups)

        self.avgpool = GlobalAvgPool()
        self.fc = Linear(cs[-1], num_classes)

    def _make_layer(self, in_channels, out_channels, blocks, stride, groups):
        layers = [Bottleneck(in_channels, out_channels, stride=stride, groups=groups)]
        for i in range(1, blocks):
            layers.append(
                Bottleneck(out_channels, out_channels, stride=1, groups=groups))
        return Sequential(*layers)

    def forward(self, x):
        x = self.stem(x)
        x = self.pool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = self.fc(x)
        return x


class Shortcut(Sequential):

    def __init__(self, in_channels, out_channels, stride):
        super().__init__()
        if stride != 1:
            self.pool = Pool2d(2, 2, type='avg', ceil_mode=True)
        if in_channels != out_channels:
            self.conv = Conv2d(in_channels, out_channels, 1, bias=False)
            self.bn = BN(out_channels)
        else:
            self.identity = Identity()


class Bottleneck(Layer):

    def __init__(self, in_channels, out_channels, stride, groups):
        super().__init__()
        channels = out_channels // 2 if groups == 32 else out_channels
        self.branch1 = Shortcut(in_channels, out_channels, stride=stride)

        self.branch2a = Conv2d(in_channels, channels, 1, bn=True, act='relu')
        self.branch2b = Conv2d(channels, channels, 3, stride=stride, groups=groups,
                               bn=True, act='relu')
        self.branch2c = Conv2d(channels, out_channels, 1,
                               bn=True)

    def forward(self, x):
        x1 = self.branch1(x)
        x2 = self.branch2a(x)
        x2 = self.branch2b(x2)
        x2 = self.branch2c(x2)
        return L.elementwise_add(x1, x2, act='relu')


def ResNeXt50_vd_64x4d():
    model = ResNeXt(layers=50, groups=64)
    return model


def ResNeXt50_vd_32x4d():
    model = ResNeXt(layers=50, groups=32)
    return model


def ResNeXt101_vd_64x4d():
    model = ResNeXt(layers=101, groups=64)
    return model


def ResNeXt101_vd_32x4d():
    model = ResNeXt(layers=101, groups=32)
    return model


def ResNeXt152_vd_64x4d():
    model = ResNeXt(layers=152, groups=64)
    return model


def ResNeXt152_vd_32x4d():
    model = ResNeXt(layers=152, groups=32)
    return model


def transpile(sd):

    import re
    # from paddle.fluid.dygraph import load_dygraph
    # sd = load_dygraph("/Users/hrvvi/Downloads/ResNeXt50_vd_32x4d_output/ResNeXt50_vd_32x4d.pdparams")[0]
    # sd = load_dygraph("/Users/hrvvi/Downloads/ResNeXt101_vd_64x4d_output/ResNeXt101_vd_64x4d.pdparams")[0]

    def map_bn_name(p):
        if p in ['mean', 'variance']:
            p = "_" + p
        elif p == 'scale':
            p = 'weight'
        elif p == 'offset':
            p = 'bias'
        return p

    new_sd = {}
    for k, v in sd.items():
        match = re.search(r"(conv|bnv)1_(\d)_(scale|offset|mean|variance|weights)", k)
        if match:
            i = int(match.group(2)) - 1
            if match.group(1) == 'conv':
                new_k = f"stem.{i}.conv.weight"
            else:
                p = map_bn_name(match.group(3))
                new_k = f"stem.{i}.bn.{p}"
            new_sd[new_k] = v
            continue

        match = re.search(r"(res|bn)([235])([a-z])_(branch[12][abc]?)_(scale|offset|mean|variance|weights)", k)
        if match:
            i1 = int(match.group(2)) - 1
            i2 = ord(match.group(3)) - ord('a')
            if match.group(1) == 'res':
                p = match.group(4)
                new_k = f"layer{i1}.{i2}.{p}.conv.weight"
            else:
                p1 = match.group(4)
                p2 = map_bn_name(match.group(5))
                new_k = f"layer{i1}.{i2}.{p1}.bn.{p2}"
            new_sd[new_k] = v
            continue

        if any(name.startswith('res4c') for name in sd.keys()):
            match = re.search(r"(res|bn)([4])([a-z])_(branch[12][abc]?)_(scale|offset|mean|variance|weights)", k)
            if match:
                i1 = int(match.group(2)) - 1
                i2 = ord(match.group(3)) - ord('a')
                if match.group(1) == 'res':
                    p = match.group(4)
                    new_k = f"layer{i1}.{i2}.{p}.conv.weight"
                else:
                    p1 = match.group(4)
                    p2 = map_bn_name(match.group(5))
                    new_k = f"layer{i1}.{i2}.{p1}.bn.{p2}"
                new_sd[new_k] = v
                continue

        else:
            match = re.search(r"(res|bn)4([ab])(\d+)?_(branch[12][abc]?)_(scale|offset|mean|variance|weights)", k)
            if match:
                if match.group(2) == 'a':
                    i = 0
                else:
                    i = int(match.group(3))
                if match.group(1) == 'res':
                    p = match.group(4)
                    new_k = f"layer3.{i}.{p}.conv.weight"
                else:
                    p1 = match.group(4)
                    p2 = map_bn_name(match.group(5))
                    new_k = f"layer3.{i}.{p1}.bn.{p2}"
                new_sd[new_k] = v
                continue


        match = re.search(r"fc_(weights|offset)", k)
        if match:
            if match.group(1) == 'weights':
                new_k = "fc.weight"
            else:
                new_k = 'fc.bias'
            new_sd[new_k] = v
            continue

        print(k)
    return new_sd